from information_extraction import UIEPredictor
from config.config_lab import UIE_ADDR


class UIEUtil:
    def __init__(self):
        self.ner_schema = ['设备', '部件', '检查内容', '故障现象', '故障原因', '排故流程', '维修方案']
        self.re_schema = {'设备': ['包含'], '部件': ['包括', '产生'], '检查内容': ['涉及'], '故障现象': ['排故'],
                          '故障原因': ['导致', '相关', '建议'], '排故流程': ['对应'], '维修方案': []}
        self.task_path = UIE_ADDR
        self.ie = None

    def predict(self, sentence):
        return self.ie(sentence)

    def clear_ie_instance(self):
        self.ie = None

    def create_ie_instance(self):
        self.ie = UIEPredictor(model='uie-base', task_path=self.task_path, schema=self.re_schema)


if __name__ == '__main__':
    uie = UIEUtil()
    uie.create_ie_instance()
    res = uie.predict(
        "最新型号的皮带机，部件包括螺丝、垫片，齿轮、皮带、锤子等。")
    print(res)
